import pygame
from pygame import QUIT, MOUSEBUTTONDOWN, K_UP, K_DOWN, K_LEFT, K_RIGHT, K_ESCAPE
from pygame import Rect
import random

ROW_COUNT = 6
COLUMN_COUNT = 6

SCREEN_SIZE = (320, 240)
SCREEN_BG_COLOR = (0, 0, 0)
SQUARE_COLOR = (255, 255, 255)

JEWEL_SIZE = 20
JEWEL_ICON_SIZE = 8
SQUARE_WIDTH = 20
LINE_JEWEL_COUNT = 3

GAME_REFRESH_TIME = 500

# const
class GAME_STATE:
	INITIAL = 0
	NORMAL = 1
	DESTROY = 2
	FILL = 3
	SWAP = 4
	
class JEWEL_TYPE:
	EMPTY = 0
	RED = 1
	BLUE = 2
	GREEN = 3
	PURPLE = 4
	
JEWEL_TYPE_MAP = dict([(key, value) for key, value in vars(JEWEL_TYPE).items() if ("__" not in key) and (key != "EMPTY")])
	
class Jewel:
	def __init__(self):
		self.type = JEWEL_TYPE.EMPTY
		self.selected = False
		self.destroy = False

class GamePlay:
	def __init__(self):		
		self.map = []		
		self.state = GAME_STATE.INITIAL
	
	# for debug
	def showMap(self):
		for j in xrange(COLUMN_COUNT):
			x = ""
			for i in xrange(ROW_COUNT):
				if self.map[j][i].destroy:
					x = x + str(self.map[j][i].type) + ". "
				else:
					x = x + str(self.map[j][i].type) + "  "
			print x
		print " "
		
	def gameReset(self):
		self.state = GAME_STATE.INITIAL
		self.initMap()
		self.fillMap()
	
	def initMap(self):
		self.map = []
		
		for j in xrange(COLUMN_COUNT):
			row = []
			for i in xrange(ROW_COUNT):
				jewel = Jewel()
				row.append(jewel)
				
			self.map.append(row)
			
	def fillMap(self):
		for j in xrange(COLUMN_COUNT):
			for i in xrange(ROW_COUNT):
				jewel = self.map[j][i]
				if jewel.type == JEWEL_TYPE.EMPTY:
					type_name = random.choice(JEWEL_TYPE_MAP.keys())
					jewel.type = JEWEL_TYPE_MAP[type_name]
					jewel.selected = False
					jewel.destroy = False
	
	def markRowStreak(self, row, col):
		if row + LINE_JEWEL_COUNT > ROW_COUNT:
			return
		
		# check if the jewel is same type in a row (default in 3 streak)
		cur_pos_type = self.map[col][row].type
		chk_row_jewel = True
		for i in xrange(LINE_JEWEL_COUNT):
			des_pos_type = self.map[col][row+i].type
			if cur_pos_type == des_pos_type:
				pass
			else:
				chk_row_jewel = False
				break
				
		# mark jewel which will be earsed
		if chk_row_jewel:
			for i in xrange(LINE_JEWEL_COUNT):
				self.map[col][row+i].destroy = True
				
	def markColumnStreak(self, row, col):
		if col + LINE_JEWEL_COUNT > COLUMN_COUNT:
			return
	
		# check if the jewel is same type in a column (default in 3 streak)
		cur_pos_type = self.map[col][row].type
		chk_col_jewel = True
		for i in xrange(LINE_JEWEL_COUNT):
			des_pos_type = self.map[col+i][row].type
			if cur_pos_type == des_pos_type:
				pass
			else:
				chk_col_jewel = False
				break
		
		# mark jewel which will be earsed
		if chk_col_jewel:
			for i in xrange(LINE_JEWEL_COUNT):
				self.map[col+i][row].destroy = True
				
	# remove a jewel
	def earseJewel(self, x, y):
		self.map[y][x].type = JEWEL_TYPE.EMPTY
		self.map[y][x].selected = False
		self.map[y][x].destroy = False
				
	# mark all jewels which is in a line streak
	# True : destroy something
	# False : destroy nothing
	def destroyJewel(self):
		self.markJewelStreak()
		
		destroy_anything = False
		for j in xrange(COLUMN_COUNT):
			for i in xrange(ROW_COUNT):
				if self.map[j][i].destroy == True:
					destroy_anything = True
					self.earseJewel(i, j)
					
		return destroy_anything
					
	def swapJewel(self, source_pos, dest_pos):
		if not (source_pos[0] == dest_pos[0] or source_pos[1] == dest_pos[1]):
			return
		
		temp = Jewel()
		source_jewel = self.map[source_pos[1]][source_pos[0]]
		dest_jewel = self.map[dest_pos[1]][dest_pos[0]]
		
		temp.type = source_jewel.type
		temp.selected = source_jewel.selected
		temp.destroy = source_jewel.destroy

		source_jewel.type = dest_jewel.type
		source_jewel.selected = False
		source_jewel.destroy = dest_jewel.destroy
		
		dest_jewel.type = temp.type
		dest_jewel.selected = False
		dest_jewel.destroy = temp.destroy
		
	def dropJewel(self, x, y):
		# already at bottom
		if y >= COLUMN_COUNT - 1:
			return
			
		source_jewel = self.map[y][x]
		dest_jewel = self.map[y+1][x]
			
		# no jewel
		if source_jewel.type == JEWEL_TYPE.EMPTY:
			return
			
		# on a non-empty jewel
		if dest_jewel.type != JEWEL_TYPE.EMPTY:
			return
			
		self.swapJewel((x, y), (x, y+1))
		
		# drop to next position
		self.dropJewel(x, y+1)
					
	def dropAllJewel(self):
		for j in xrange(COLUMN_COUNT - 1, -1, -1):
			for i in xrange(ROW_COUNT):
				self.dropJewel(i, j)
				
	def markJewelStreak(self):
		for j in xrange(COLUMN_COUNT):
			for i in xrange(ROW_COUNT):
				self.markRowStreak(i, j)
				self.markColumnStreak(i, j)
				
	def selectJewel(self, x, y):
		print x, y
		if self.state != GAME_STATE.NORMAL:
			return
			
		self.map[y][x].selected = True
		
		selected_jewel_around = False
		swap_pos = (-1, -1)
		if x - 1 >= 0 and self.map[y][x-1].selected:
			selected_jewel_around = True
			swap_pos = (x - 1, y)
		elif x + 1 < ROW_COUNT and self.map[y][x+1].selected:
			selected_jewel_around = True
			swap_pos = (x + 1, y)
		elif y - 1 >= 0 and self.map[y-1][x].selected:
			selected_jewel_around = True
			swap_pos = (x, y - 1)
		elif y + 1 < COLUMN_COUNT and self.map[y+1][x].selected:
			selected_jewel_around = True
			swap_pos = (x, y + 1)
		
		if selected_jewel_around:
			self.swapJewel((x, y), swap_pos)
			self.state = GAME_STATE.SWAP
	
	def update(self):
		if self.state == GAME_STATE.INITIAL:
			self.gameReset()
			self.state = GAME_STATE.NORMAL
			
		elif self.state == GAME_STATE.NORMAL:
			destroyed = False
			destroyed = self.destroyJewel()
			
			if destroyed:
				self.state = GAME_STATE.DESTROY
			
		elif self.state == GAME_STATE.DESTROY:
			self.dropAllJewel()
			self.state = GAME_STATE.FILL
			
		elif self.state == GAME_STATE.FILL:
			self.fillMap()
			self.state = GAME_STATE.NORMAL
			
		elif self.state == GAME_STATE.SWAP:
			self.state = GAME_STATE.NORMAL
			
		#self.showMap()
		
			
		
class GameViewer:
	def __init__(self):
		pygame.init()
		
		self.screen = pygame.display.set_mode(SCREEN_SIZE)
		
	def draw(self, map):
		# fill background
		self.screen.fill(SCREEN_BG_COLOR)
		
		# draw map
		for y, row in enumerate(map):
			for x, jewel in enumerate(row):
				self.drawJewel((x, y), jewel.type)
					
				if jewel.selected == True:
					self.drawSquare((x, y))
		
		pygame.display.flip()
		
	def drawJewel(self, map_pos, jewel_type):
		x = map_pos[0] * JEWEL_SIZE + 2
		y = map_pos[1] * JEWEL_SIZE + 2
		rect = Rect(x, y, JEWEL_ICON_SIZE*2, JEWEL_ICON_SIZE*2)
		
		jewel_color = (0, 0, 0)
		if jewel_type == JEWEL_TYPE.RED:
			jewel_color = (255, 0, 0)
		elif jewel_type == JEWEL_TYPE.BLUE:
			jewel_color = (0, 0, 255)
		elif jewel_type == JEWEL_TYPE.GREEN:
			jewel_color = (0, 255, 0)
		elif jewel_type == JEWEL_TYPE.PURPLE:
			jewel_color = (255, 0, 255)
		
		pygame.draw.rect(self.screen, jewel_color, rect, 0)

	def drawSquare(self, map_pos):
		x = map_pos[0] * JEWEL_SIZE
		y = map_pos[1] * JEWEL_SIZE
		rect = Rect(x, y, SQUARE_WIDTH, SQUARE_WIDTH)
		pygame.draw.rect(self.screen, SQUARE_COLOR, rect, 1)

def PosToPixel(pos_x, pos_y):
	x = pos_x * JEWEL_SIZE
	y = pos_y * JEWEL_SIZE
	
	return x, y
	
def PixelToPos(x, y):
	pos_x = x / JEWEL_SIZE
	pos_y = y / JEWEL_SIZE
	
	return pos_x, pos_y
	
	
		
class GameClient:
	def __init__(self):
		self.view = GameViewer()
		self.game_play = GamePlay()
	
	def run(self):		
		while True:
			self.processInput()
			self.game_play.update()
			self.view.draw(self.game_play.map)
			pygame.time.delay(GAME_REFRESH_TIME)
			
	def processInput(self):
		for event in pygame.event.get():
			if event.type == QUIT:
				exit()
			elif event.type == MOUSEBUTTONDOWN:
				x, y = event.pos
				map_x, map_y = PixelToPos(x, y)
				self.game_play.selectJewel(map_x, map_y)
				
				

if __name__ == "__main__":
	game = GameClient()
	game.run()